from comm_env import NetworkEnv
from comm_env.comm_env2 import RETURN_T, OBS_T
import gym
from typing import List, Dict


class SchemeEnvWrapper(gym.Wrapper):
    def __init__(self, env: NetworkEnv, scheme_steps: List[int],
                 parameters_list: List[dict]):
        """使得环境按照`scheme_steps`中的时间点更改环境，`parameters_list`为环境初始化参数。
        另外，不支持被gym.Wrapper已经包装过的环境
        env: 被包装的gym.Env
        scheme_steps: 到达指定step后按照对应参数初始化环境, 没运行一次则step+1
            e.g. [1000, 2000, 3500],
            1000步时初始化为`parameters_list[1]`
        parameters_list: 参数字典，要求对应`env`的初始化参数
        """
        super().__init__(env)
        assert len(scheme_steps) == len(parameters_list)
        self.scheme_steps = scheme_steps
        self.parameters_list = parameters_list
        self._stage = 0  # 下一阶段指针, 用以指定初始化的参数
        self._total_step_counter = 0
        self._trans_signal = False

    def step(self, act) -> RETURN_T:
        self._total_step_counter += 1
        obs_, rew, done, info = self.env.step(act)
        if self._total_step_counter == self.scheme_steps[self._stage]:
            done = True
            self._trans_signal = True
        return obs_, rew, done, info

    def reset(self) -> OBS_T:
        if self._trans_signal:
            env = NetworkEnv(**self.parameters_list[self._stage])
            self._trans_signal = False
            self._stage += 1
            self.env = env
        return self.env.reset()


class SchemeProbEnvWrapper(gym.Wrapper):
    def __init__(
        self,
        env: NetworkEnv,
        scheme_link_prob: List[List[float]],
        scheme_eng_prob: List[List[float]],
        scheme_steps: List[int],
    ):
        super().__init__(env)
        assert len(scheme_eng_prob) == len(scheme_link_prob)
        assert len(scheme_eng_prob) == len(scheme_steps)
        self.scheme_link_prob = scheme_link_prob
        self.scheme_eng_prob = scheme_eng_prob
        self.scheme_steps = scheme_steps
        self.scheme_steps.append(float('inf'))
        self._next_stage = 0
        self._global_step_counter = 0
        self._trans_signal = False

    def reset(self) -> OBS_T:
        if self._trans_signal:
            link_probs = self.scheme_link_prob[self._next_stage]
            eng_probs = self.scheme_eng_prob[self._next_stage]
            self.env.set_link_prob(link_probs)
            self.env.set_eng_prob(eng_probs)
            self._next_stage += 1
            self._trans_signal = False
        return self.env.reset()

    def step(self, act) -> RETURN_T:
        obs_, rew, done, info = self.env.step(act)
        self._global_step_counter += 1
        if self._global_step_counter >= self.scheme_steps[self._next_stage]:
            done = True
            self._trans_signal = True
        return obs_, rew, done, info
